#!/bin/bash

# AIQuant 系统停止脚本
# 用于停止AIQuant量化投研系统的所有服务

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 停止进程函数
stop_process() {
    local pid_file=$1
    local service_name=$2
    
    if [ -f "$pid_file" ]; then
        local pid=$(cat "$pid_file")
        if ps -p $pid > /dev/null 2>&1; then
            log_info "停止 $service_name (PID: $pid)..."
            kill $pid
            
            # 等待进程结束
            local count=0
            while ps -p $pid > /dev/null 2>&1 && [ $count -lt 10 ]; do
                sleep 1
                count=$((count + 1))
            done
            
            # 如果进程仍在运行，强制终止
            if ps -p $pid > /dev/null 2>&1; then
                log_warning "强制终止 $service_name..."
                kill -9 $pid
            fi
            
            log_success "$service_name 已停止"
        else
            log_warning "$service_name 进程不存在 (PID: $pid)"
        fi
        rm -f "$pid_file"
    else
        log_warning "$service_name PID文件不存在"
    fi
}

# 停止API服务
stop_api() {
    log_info "停止API服务..."
    stop_process ".api.pid" "API服务"
    
    # 额外检查并终止可能的uvicorn进程
    pkill -f "uvicorn.*aiquant.api.main:app" 2>/dev/null || true
}

# 停止Web服务
stop_web() {
    log_info "停止Web服务..."
    stop_process ".web.pid" "Web服务"
    
    # 额外检查并终止可能的streamlit进程
    pkill -f "streamlit.*aiquant/web/dashboard.py" 2>/dev/null || true
}

# 停止后台任务
stop_background_tasks() {
    log_info "停止后台任务..."
    
    # 停止数据收集器
    stop_process ".data_collector.pid" "数据收集器"
    
    # 停止市场监控器
    stop_process ".market_monitor.pid" "市场监控器"
    
    # 停止异常检测器
    stop_process ".anomaly_detector.pid" "异常检测器"
    
    # 额外检查并终止可能的后台进程
    pkill -f "aiquant.market_observation" 2>/dev/null || true
}

# 停止Redis（可选）
stop_redis() {
    if [ "$1" = "--with-redis" ]; then
        log_info "停止Redis服务..."
        if pgrep -x "redis-server" > /dev/null; then
            pkill -x redis-server
            log_success "Redis服务已停止"
        else
            log_warning "Redis服务未运行"
        fi
    fi
}

# 清理临时文件
cleanup() {
    log_info "清理临时文件..."
    
    # 清理PID文件
    rm -f .api.pid .web.pid .data_collector.pid .market_monitor.pid .anomaly_detector.pid
    
    # 清理临时缓存
    find . -name "*.pyc" -delete 2>/dev/null || true
    find . -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true
    
    log_success "临时文件清理完成"
}

# 显示停止状态
show_status() {
    log_info "检查服务状态..."
    
    local running_services=0
    
    # 检查API服务
    if lsof -Pi :8000 -sTCP:LISTEN -t >/dev/null 2>&1; then
        echo "API服务仍在运行 (端口8000)"
        running_services=$((running_services + 1))
    fi
    
    # 检查Web服务
    if lsof -Pi :8501 -sTCP:LISTEN -t >/dev/null 2>&1; then
        echo "Web服务仍在运行 (端口8501)"
        running_services=$((running_services + 1))
    fi
    
    # 检查后台进程
    if pgrep -f "aiquant" > /dev/null 2>&1; then
        echo "后台进程仍在运行:"
        pgrep -f "aiquant" | while read pid; do
            ps -p $pid -o pid,cmd --no-headers 2>/dev/null || true
        done
        running_services=$((running_services + 1))
    fi
    
    if [ $running_services -eq 0 ]; then
        log_success "所有服务已停止"
    else
        log_warning "仍有 $running_services 个服务在运行"
    fi
}

# 强制停止所有相关进程
force_stop() {
    log_warning "强制停止所有AIQuant相关进程..."
    
    # 强制终止所有相关进程
    pkill -f "aiquant" 2>/dev/null || true
    pkill -f "uvicorn.*aiquant" 2>/dev/null || true
    pkill -f "streamlit.*aiquant" 2>/dev/null || true
    
    # 等待进程终止
    sleep 2
    
    # 检查是否还有进程运行
    if pgrep -f "aiquant" > /dev/null 2>&1; then
        log_warning "使用SIGKILL强制终止剩余进程..."
        pkill -9 -f "aiquant" 2>/dev/null || true
    fi
    
    log_success "强制停止完成"
}

# 主函数
main() {
    log_info "停止AIQuant量化投研系统..."
    
    case "${1:-}" in
        --force|-f)
            force_stop
            cleanup
            ;;
        --with-redis)
            stop_api
            stop_web
            stop_background_tasks
            stop_redis --with-redis
            cleanup
            ;;
        --help|-h)
            echo "AIQuant系统停止脚本"
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --help, -h        显示帮助信息"
            echo "  --force, -f       强制停止所有进程"
            echo "  --with-redis      同时停止Redis服务"
            echo "  --status          仅显示服务状态"
            exit 0
            ;;
        --status)
            show_status
            exit 0
            ;;
        "")
            stop_api
            stop_web
            stop_background_tasks
            cleanup
            ;;
        *)
            log_error "未知参数: $1"
            echo "使用 --help 查看帮助信息"
            exit 1
            ;;
    esac
    
    show_status
    log_success "AIQuant系统停止完成!"
}

# 执行主函数
main "$@"